// 维护哪些流hash已经被推理，哪些应该被推理
#include <iostream>
#include <unordered_map>
#include <vector>
#include <queue>
#include <chrono>
#include <sstream>
#include <headers/SocketServer.h>

#define MAX_SEND_ONCE 2000

// 定义日志级别
enum class LogLevel {
    INFO,
    WARNING,
    ERROR
};

// 定义表中的每个条目的结构体
struct HashEntry {
    uint8_t packetCount : 7;  // 包数，最大127，足够记录到100
    bool isInferred : 1;      // 是否已经被推理
    
    HashEntry() : packetCount(0), isInferred(false) {}
};

// 维护哈希表的类
class HashTableManager {
private:
    static constexpr size_t MAX_HASH_TABLE_SIZE = 1000000;  // 哈希表最大容量
    static constexpr int RECONNECT_DELAY_MS = 5000;      // 重连延迟（毫秒）
    static constexpr int PACKET_THRESHOLD = 100;         // 包数阈值
    
    std::unordered_map<uint32_t, HashEntry> hashTable;      // 存储哈希表数据
    std::queue<uint32_t> inferenceQueue;                    // 待推理队列
    SocketServer server;
    std::thread server_thread;                              // 服务器线程
    std::atomic<bool> running;                              // 控制服务器线程运行
    std::mutex hashTable_mutex;                             // 保护哈希表的互斥锁
    std::mutex queue_mutex;                                 // 保护待推理队列的互斥锁

    // 记录日志
    void log(LogLevel level, const std::string& message) {
        static const char* level_strings[] = {"INFO", "WARNING", "ERROR"};
        std::stringstream ss;
        auto now = std::chrono::system_clock::now();
        auto now_c = std::chrono::system_clock::to_time_t(now);
        ss << std::put_time(std::localtime(&now_c), "%Y-%m-%d %H:%M:%S")
           << " [" << level_strings[static_cast<int>(level)] << "] "
           << message << std::endl;
        std::cout << ss.str();
        std::cout.flush();
    }

    // 服务器线程函数
    void server_routine() {
        std::cerr << "server_routine线程启动" << std::endl;
        log(LogLevel::INFO, "开始初始化服务器...");
        // 初始化服务器
        if (!server.init()) {
            log(LogLevel::ERROR, "服务器初始化失败");
            throw std::runtime_error("服务器初始化失败");
        }

        // 等待客户端连接
        log(LogLevel::INFO, "等待Python客户端连接...");
        
        // 阻塞，等待连接
        server.accept_connection();
        log(LogLevel::INFO, "Python客户端连接成功");

        // 连接成功后，每12秒发送一次数据
        while (running) {
            try {
                talk_to_python();
                std::this_thread::sleep_for(std::chrono::seconds(12));
            } catch (const std::exception& e) {
                log(LogLevel::ERROR, std::string("与Python客户端通信异常: ") + e.what());
                break;  // 跳出内层循环，重新连接
            }
        }
        log(LogLevel::INFO, "客户端连接断开，等待重新连接...");
    }

public:
    // 构造函数
    HashTableManager() : running(true) {
        // 启动服务器线程
        server_thread = std::thread(&HashTableManager::server_routine, this);
    }

    // 析构函数
    ~HashTableManager() {
        // 停止服务器线程并关闭 socket，使 accept 能被中断
        running = false;
        if (server_thread.joinable()) {
            server_thread.join();
        }
    }

    // 更新指定哈希值的包数,返回会话数
    long long updatePacketCount(uint32_t hashVal) {
        std::lock_guard<std::mutex> table_lock(hashTable_mutex);
        
        // 检查哈希表大小是否超过限制
        if (hashTable.size() >= MAX_HASH_TABLE_SIZE && hashTable.find(hashVal) == hashTable.end()) {
            log(LogLevel::WARNING, "哈希表已达到最大容量限制");
            return static_cast<long long>(MAX_HASH_TABLE_SIZE);
        }
        // std::cout<<"called"<<std::endl;
        auto it = hashTable.find(hashVal);
        if (it == hashTable.end()) {
            // std::cout<<"not found"<<std::endl;
            hashTable[hashVal] = HashEntry();
            hashTable[hashVal].packetCount = 1;
        } 
        else if (it->second.packetCount < PACKET_THRESHOLD && !it->second.isInferred) {
            // std::cout<<"hashVal:"<<hashVal<<", packetCount:"<<it->second.packetCount<<std::endl;
            ++(it->second.packetCount);
            // std::cout<<"added"<<std::endl;
            // 如果达到阈值，加入待推理队列
            if (it->second.packetCount >= PACKET_THRESHOLD) {
                std::lock_guard<std::mutex> queue_lock(queue_mutex);
                inferenceQueue.push(hashVal);
                // log(LogLevel::INFO, "Hash " + std::to_string(hashVal) + " 达到推理阈值");
            }
        }
        return static_cast<long long>(hashTable.size());
    }

    // 获取待推理的哈希值，最多1000个
    std::vector<uint32_t> getHashesToInfer() {
        std::vector<uint32_t> result;
        std::lock_guard<std::mutex> queue_lock(queue_mutex);
        
        // 计算要取出的元素数量
        size_t count = std::min(inferenceQueue.size(), static_cast<size_t>(MAX_SEND_ONCE));
        result.reserve(count);  // 预分配空间
        
        // 一次性取出指定数量的元素
        for (size_t i = 0; i < count; ++i) {
            result.push_back(inferenceQueue.front());
            inferenceQueue.pop();
        }
        
        return result;
    }

    // 标记某个哈希值为已推理
    void markAsInferred(uint32_t hashVal) {
        std::lock_guard<std::mutex> lock(hashTable_mutex);
        auto it = hashTable.find(hashVal);
        if (it != hashTable.end()) {
            it->second.isInferred = true;
        }
    }

    void talk_to_python() {
        auto hashesToBeInferred = getHashesToInfer();
        if (hashesToBeInferred.empty()) {
            // log(LogLevel::INFO, "没有需要推理的数据");
            return; // 如果没有需要推理的数据，直接返回
        }

        bool isSendSuccess = server.send_vector(hashesToBeInferred);
        if(isSendSuccess) {
            for(auto hash : hashesToBeInferred) {
                markAsInferred(hash);
            }
            log(LogLevel::INFO, "成功发送 " + std::to_string(hashesToBeInferred.size()) + " 个哈希值到Python客户端");
        } else {
            log(LogLevel::ERROR, "发送失败");
            throw std::runtime_error("与Python客户端通信失败");
        }
    }

    // // 计算哈希表中元素的数量（hash 数）
    // size_t getHashCount() const {
    //     std::lock_guard<std::mutex> lock(hashTable_mutex);
    //     return hashTable.size();
    // }

    
    // // 估算整体内存占用
    // size_t getEstimatedMemoryUsage() const {
    //     // 一个键（uint32_t）的大小
    //     size_t keySize = sizeof(uint32_t);
    //     // 一个值（HashEntry）的大小
    //     size_t valueSize = sizeof(HashEntry);
    //     // 每个元素的大小
    //     size_t elementSize = keySize + valueSize;
    //     // 哈希表中元素的数量
    //     size_t elementCount = hashTable.size();
    //     // 整体内存占用
    //     return elementCount * elementSize;
    // }
};